﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using System.Windows.Forms;
using Microsoft.Win32;
using System.ComponentModel;
using System.IO;

namespace RoketPack
{
    class Protocol
    {
        internal static readonly string m_ProtocolHandlerDirectory =
            Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles), "RoketPack");
        internal static readonly string m_ProtocolHandlerFile =
            Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles), "RoketPack\\RoketProtocol.exe");
        internal static readonly string m_ProtocolHandlerDLL =
            Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles), "RoketPack\\RoketPackAPI.dll");

        /// <summary>
        /// Called at program start, this function handles situations where RoketPack
        /// should be associating itself with the roket: protocol while it has elevated
        /// privileges.  Returns whether the application should immediately exit.
        /// </summary>
        /// <param name="args">The command line arguments passed to the program.</param>
        /// <returns>Whether the application should immediately exit.</returns>
        internal static bool HandleAssociation(string[] args)
        {
            // Check to see whether we've been asked to associate the roket:
            // protocol with the current executable.
            if (args.Length >= 1 && args[0] == "--associate")
            {
                // Make sure that we run as administrator when we're
                // doing this.
                if (!UAC.IsAdmin())
                {
                    // Launch ourselves as administrator.
                    ProcessStartInfo info = new ProcessStartInfo();
                    info.FileName = UNC.UniversalApplicationPath;
                    info.UseShellExecute = true;
                    info.Verb = "runas"; // Provides Run as Administrator
                    info.Arguments = "--associate";

                    // Check to see if launch failed.
                    if (Process.Start(info) == null)
                        Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);

                    return true;
                }
                else
                {
                    try
                    {
                        // Attempt to copy our version of RoketProtocol.exe into Program Files
                        // before we set the registry key.
                        if (File.Exists(Path.Combine(Environment.CurrentDirectory, "RoketProtocol.exe")))
                        {
                            if (!Directory.Exists(Protocol.m_ProtocolHandlerDirectory))
                            {
                                try
                                {
                                    Directory.CreateDirectory(Protocol.m_ProtocolHandlerDirectory);
                                }
                                catch (UnauthorizedAccessException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                                catch (IOException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                            }
                            if (!File.Exists(Protocol.m_ProtocolHandlerFile))
                            {
                                try
                                {
                                    File.Copy(Path.Combine(Environment.CurrentDirectory, "RoketProtocol.exe"), Protocol.m_ProtocolHandlerFile, true);
                                }
                                catch (UnauthorizedAccessException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                                catch (IOException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                            }
                            if (!File.Exists(Protocol.m_ProtocolHandlerDLL))
                            {
                                try
                                {
                                    File.Copy(Path.Combine(Environment.CurrentDirectory, "RoketPackAPI.dll"), Protocol.m_ProtocolHandlerDLL, true);
                                }
                                catch (UnauthorizedAccessException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                                catch (IOException)
                                {
                                    Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                                    return true;
                                }
                            }
                        }

                        // Set up the main registry key for handling the roket: protocol.
                        RegistryKey r = Registry.ClassesRoot;
                        r = r.CreateSubKey("roket", RegistryKeyPermissionCheck.ReadWriteSubTree);
                        r.SetValue("", "URL: RoketPack Installation Protocol");
                        r.SetValue("URL Protocol", "");
                        r.CreateSubKey("DefaultIcon").SetValue("", Protocol.m_ProtocolHandlerFile);
                        RegistryKey c = r.CreateSubKey("shell").CreateSubKey("open").CreateSubKey("command");
                        c.SetValue("", Protocol.m_ProtocolHandlerFile + " \"%1\"");
                        c.Close();
                        r.Close();
                    }
                    catch (Exception ex)
                    {
                        Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                    }
                    return true;
                }
            }

            return false;
        }

        /// <summary>
        /// Ensures that the roket: protocol is associated with RoketPack.  Returns
        /// true if the application was successfully associated.
        /// </summary>
        /// <returns>If the application was successfully associated.</returns>
    internal static bool EnsureAssociation()
    {
        // Check to make sure RoketPack is associated.
        if (!Protocol.IsAssociated())
        {
            // If we are already running as administrator, we can
            // just call HandleAssociation.  If we didn't do this,
            // then Process.Start would attempt to re-use the existing
            // process (since it's under the same user) and we wouldn't
            // be able to wait until it's finished.
            if (UAC.IsAdmin())
            {
                Protocol.HandleAssociation(new string[] { "--associate" });
                return Protocol.IsAssociated();
            }

            ProcessStartInfo info = new ProcessStartInfo();
            info.FileName = UNC.UniversalApplicationPath;
            info.UseShellExecute = true;
            if (!UAC.IsAdmin())
                info.Verb = "runas"; // Provides Run as Administrator
            info.Arguments = "--associate";
            Process proc = null;
            try
            {
                proc = Process.Start(info);
            }
            catch (Win32Exception)
            {
                Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                return false;
            }

            if (null != proc)
            {
                // Wait until the association is complete.
                proc.WaitForExit();
                return Protocol.IsAssociated();
            }
            else
            {
                Errors.Raise(Errors.ErrorType.ERROR_CAN_NOT_ASSOCIATE_PROTOCOL);
                return false;
            }
        }
        else
            return true;
    }

        /// <summary>
        /// Returns whether or not RoketPack is associated with the roket: protocol.
        /// </summary>
        /// <returns>Whether or not RoketPack is associated with the roket: protocol.</returns>
        internal static bool IsAssociated()
        {
            // Check to make sure RoketProtocol is installed.
            if (!File.Exists(Protocol.m_ProtocolHandlerFile))
                return false;
            if (!File.Exists(Protocol.m_ProtocolHandlerDLL))
                return false;

            // Check registry is in order.
            RegistryKey r = Registry.ClassesRoot;
            bool notassociated = false;
            try
            {
                r = r.OpenSubKey("roket", false);
                if ((string)r.GetValue("", "<nonexistant>") != "URL: RoketPack Installation Protocol" ||
                    (string)r.GetValue("URL Protocol", "<nonexistant>") == "<nonexistant>")
                    notassociated = true;
                if ((string)r.OpenSubKey("DefaultIcon").GetValue("", "<nonexistant>") != Protocol.m_ProtocolHandlerFile)
                    notassociated = true;
                RegistryKey c = r.OpenSubKey("shell").OpenSubKey("open").OpenSubKey("command");
                if ((string)c.GetValue("", "<nonexistant>") != Protocol.m_ProtocolHandlerFile + " \"%1\"")
                    notassociated = true;
            }
            catch (ApplicationException)
            {
                // TODO: Find out what kind of exceptions will be raised due to access problems, etc..
                return false;
            }
            catch (NullReferenceException)
            {
                return false;
            }

            return !notassociated;
        }
    }
}
